Model/orchestration/audit_generator_orchestrator.py

79 lines
2.7 KiB
Python

from __future__ import annotations
from collections.abc import Callable
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
import openpyxl
from domain.magicplan.ventilation_audit import populate_sheet
from infrastructure.postgres.uploaded_file_table import (
FileSourceEnum,
FileTypeEnum,
UploadedFile,
)
from infrastructure.s3.s3_client import S3Client
if TYPE_CHECKING:
from orchestration.audit_generator_unit_of_work import AuditGeneratorUnitOfWork
_TEMPLATE_PATH = Path(__file__).parent.parent / "applications" / "audit_generator" / "Master Sero Template - Data Extraction.xlsx"
_SHEET_NAME = "D1 Ventilation"
def _serialise_workbook(wb: Any) -> bytes:
buf = BytesIO()
wb.save(buf)
return buf.getvalue()
class AuditGeneratorOrchestrator:
def __init__(
self,
hubspot_deal_id: str,
s3_client: S3Client,
uow_factory: Callable[[], "AuditGeneratorUnitOfWork"],
) -> None:
self._hubspot_deal_id = hubspot_deal_id
self._s3_client = s3_client
self._uow_factory = uow_factory
def run(self) -> None:
with self._uow_factory() as uow:
uploaded_file = uow.uploaded_file.get_latest_by_hubspot_deal_id(
self._hubspot_deal_id, FileTypeEnum.MAGIC_PLAN_JSON
)
if uploaded_file is None:
raise ValueError(
f"No MagicPlan JSON has been uploaded for deal {self._hubspot_deal_id!r}"
)
plan = uow.magic_plan.get_plan_by_uploaded_file_id(cast(int, uploaded_file.id))
if plan is None:
raise ValueError(
f"MagicPlan JSON exists for deal {self._hubspot_deal_id!r} "
"but the plan is not yet parsed into the database"
)
wb = openpyxl.load_workbook(_TEMPLATE_PATH)
sheet = wb[_SHEET_NAME]
populate_sheet(sheet, plan)
xlsx_bytes = _serialise_workbook(wb)
s3_key = (
f"documents/hubspot_deal_id/{self._hubspot_deal_id}/ventilation_audit.xlsx"
)
self._s3_client.put_object(s3_key, xlsx_bytes)
new_row = UploadedFile(
s3_file_bucket=self._s3_client.bucket,
s3_file_key=s3_key,
s3_upload_timestamp=datetime.now(timezone.utc),
hubspot_deal_id=self._hubspot_deal_id,
file_type=FileTypeEnum.VENTILATION_AUDIT.value,
file_source=FileSourceEnum.AUDIT_GENERATOR.value,
)
uow.uploaded_file.insert(new_row)
uow.commit()