Model/scripts/run_audit_generator_local.py
2026-06-11 09:32:14 +00:00

83 lines
2.3 KiB
Python

"""
Run audit_generator locally.
Usage:
cd /workspaces/model
python scripts/run_audit_generator_local.py [<hubspot_deal_id>]
Prompts for deal ID and S3 destination (local file or real S3) if not supplied
"""
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Any, Union
import boto3
# Load .env before importing infra modules
from dotenv import load_dotenv
load_dotenv(Path(__file__).parent.parent / "backend" / ".env")
from infrastructure.postgres.config import PostgresConfig
from infrastructure.postgres.engine import make_engine, make_session
from infrastructure.s3.s3_client import S3Client
from orchestration.audit_generator_orchestrator import AuditGeneratorOrchestrator
from orchestration.audit_generator_unit_of_work import AuditGeneratorUnitOfWork
class _LocalS3Client:
"""Writes to local filesystem instead of S3."""
def __init__(self, output_dir: Path) -> None:
self._output_dir = output_dir
self._output_dir.mkdir(parents=True, exist_ok=True)
@property
def bucket(self) -> str:
return "local"
def get_object(self, key: str) -> bytes:
raise NotImplementedError
def put_object(self, key: str, body: bytes) -> str:
dest = self._output_dir / Path(key).name
dest.write_bytes(body)
print(f"Saved: {dest}")
return str(dest)
def _make_s3_client() -> Union[S3Client, "_LocalS3Client"]:
use_real = input("Use real S3? [y/N]: ").strip().lower() == "y"
if use_real:
bucket = "retrofit-energy-assessments-dev"
boto3_client: Any = boto3.client
return S3Client(boto_s3_client=boto3_client("s3"), bucket=bucket)
output_dir = Path(__file__).parent.parent / "local_output"
return _LocalS3Client(output_dir)
def main() -> None:
deal_id = sys.argv[1] if len(sys.argv) > 1 else input("hubspot_deal_id: ").strip()
s3_client = _make_s3_client()
engine = make_engine(PostgresConfig.from_env(os.environ))
def session_factory() -> Any:
return make_session(engine)
def uow_factory() -> AuditGeneratorUnitOfWork:
return AuditGeneratorUnitOfWork(session_factory)
AuditGeneratorOrchestrator(
hubspot_deal_id=deal_id,
s3_client=s3_client, # type: ignore[arg-type]
uow_factory=uow_factory,
).run()
if __name__ == "__main__":
main()