158 lines
5.1 KiB
Python
158 lines
5.1 KiB
Python
"""Command line entry point for explicit scheduled runs and backfills."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import sys
|
|
from datetime import date
|
|
from pathlib import Path
|
|
from typing import Optional, Sequence, Tuple
|
|
|
|
from company_reports.contracts import COMPANY_NAMES, ErrorCode, RESULT_SCHEMA_VERSION
|
|
from company_reports.publishing import ArtifactToolBuilder, AtomicReportPublisher
|
|
from company_reports.repository import DatabaseConfig, PostgresReportRepository, RepositoryError
|
|
from company_reports.service import CompanyReportService, RunRequest, write_batch_result
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
MONTH_RE = re.compile(r"^(\d{4})-(\d{2})$")
|
|
|
|
|
|
class RequestFailure(ValueError):
|
|
pass
|
|
|
|
|
|
class SafeArgumentParser(argparse.ArgumentParser):
|
|
def error(self, message: str) -> None:
|
|
raise RequestFailure(message)
|
|
|
|
|
|
def _parse_month(value: str) -> Tuple[int, int]:
|
|
match = MONTH_RE.fullmatch(value)
|
|
if not match:
|
|
raise RequestFailure("month must use YYYY-MM")
|
|
year, month = int(match.group(1)), int(match.group(2))
|
|
if month < 1 or month > 12:
|
|
raise RequestFailure("month is invalid")
|
|
return year, month
|
|
|
|
|
|
def _parse_date(value: str) -> date:
|
|
try:
|
|
return date.fromisoformat(value)
|
|
except ValueError:
|
|
raise RequestFailure("as-of date must use YYYY-MM-DD") from None
|
|
|
|
|
|
def _parser() -> SafeArgumentParser:
|
|
parser = SafeArgumentParser(prog="company-reports")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
generate = subparsers.add_parser(
|
|
"generate",
|
|
help="generate or explicitly backfill one company-report month",
|
|
)
|
|
generate.add_argument("--month", required=True, help="report month in YYYY-MM")
|
|
generate.add_argument(
|
|
"--as-of",
|
|
required=True,
|
|
help="phase end: the 10th, 20th, or natural month end",
|
|
)
|
|
generate.add_argument(
|
|
"--company",
|
|
action="append",
|
|
choices=COMPANY_NAMES,
|
|
help="repeat to select companies; defaults to all five",
|
|
)
|
|
generate.add_argument(
|
|
"--output-root",
|
|
help="controlled output root inside the project",
|
|
)
|
|
generate.add_argument(
|
|
"--node-binary",
|
|
help="Node.js executable; defaults to COMPANY_REPORT_NODE_BINARY or PATH",
|
|
)
|
|
generate.add_argument(
|
|
"--artifact-tool-module",
|
|
help="absolute path to artifact_tool.mjs when package resolution is unavailable",
|
|
)
|
|
return parser
|
|
|
|
|
|
def _failure(code: str, stage: str) -> dict:
|
|
return {
|
|
"schema_version": RESULT_SCHEMA_VERSION,
|
|
"status": "failed",
|
|
"errors": [{"code": code, "stage": stage, "record_ids": []}],
|
|
}
|
|
|
|
|
|
def _emit(payload: dict) -> None:
|
|
sys.stdout.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")
|
|
|
|
|
|
def main(argv: Optional[Sequence[str]] = None) -> int:
|
|
try:
|
|
args = _parser().parse_args(argv)
|
|
year, month = _parse_month(args.month)
|
|
as_of_date = _parse_date(args.as_of)
|
|
companies = tuple(args.company or COMPANY_NAMES)
|
|
request = RunRequest(year, month, as_of_date, companies)
|
|
request.validate()
|
|
output_root = (
|
|
Path(args.output_root).expanduser()
|
|
if args.output_root
|
|
else PROJECT_ROOT / "outputs" / "company_reports"
|
|
)
|
|
if not output_root.is_absolute():
|
|
output_root = PROJECT_ROOT / output_root
|
|
output_root = output_root.resolve()
|
|
config = DatabaseConfig.from_environment()
|
|
repository = PostgresReportRepository(config)
|
|
builder = ArtifactToolBuilder(
|
|
PROJECT_ROOT / "company_reports" / "xlsx" / "build_workbook.mjs",
|
|
node_binary=args.node_binary,
|
|
artifact_tool_module=(
|
|
Path(args.artifact_tool_module).expanduser()
|
|
if args.artifact_tool_module
|
|
else None
|
|
),
|
|
)
|
|
publisher = AtomicReportPublisher(PROJECT_ROOT, output_root)
|
|
service = CompanyReportService(
|
|
repository,
|
|
builder,
|
|
publisher,
|
|
output_root / ".staging",
|
|
)
|
|
result = service.run(request)
|
|
result_path = (
|
|
output_root
|
|
/ f"{year:04d}"
|
|
/ f"{month:02d}"
|
|
/ "results"
|
|
/ f"company-reports-{year:04d}-{month:02d}-{as_of_date.isoformat()}.result.json"
|
|
)
|
|
write_batch_result(result_path, result)
|
|
_emit(result.to_dict())
|
|
return result.exit_code
|
|
except (RequestFailure, ValueError) as error:
|
|
_emit(_failure(ErrorCode.REQUEST_INVALID, "request"))
|
|
return 3
|
|
except RepositoryError as error:
|
|
code = (
|
|
ErrorCode.REQUEST_INVALID
|
|
if error.code == ErrorCode.REQUEST_INVALID
|
|
else ErrorCode.INTERNAL_ERROR
|
|
)
|
|
_emit(_failure(code, "configuration"))
|
|
return 3 if code == ErrorCode.REQUEST_INVALID else 4
|
|
except Exception:
|
|
_emit(_failure(ErrorCode.INTERNAL_ERROR, "internal"))
|
|
return 4
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|