import gzip import json from datetime import datetime, timezone from typing import Optional from datatypes.magicplan.api.response import ( MagicPlanPlan, PlanSummary, PlansListResponse, ) from datatypes.magicplan.domain.mapper import map_plan from datatypes.magicplan.domain.models import Plan from backend.app.db.connection import db_session from backend.app.db.functions.magic_plan_functions import save_plan from backend.app.db.models.uploaded_file import ( FileSourceEnum, FileTypeEnum, UploadedFile, ) from backend.magic_plan.address_matcher import find_matching_plan from backend.magic_plan.magic_plan_client import MagicPlanClient from backend.magic_plan.magic_plan_trigger_request import MagicPlanTriggerRequest from utils.logger import setup_logger from utils.s3 import save_data_to_s3 logger = setup_logger() class MagicPlanService: def __init__(self, client: MagicPlanClient, s3_bucket: str) -> None: self._client = client self._s3_bucket = s3_bucket def run(self, request: MagicPlanTriggerRequest) -> Plan: address = request.address uprn = request.uprn if uprn is not None: logger.info("MagicPlanService.run uprn=%s", uprn) plans_response: PlansListResponse = self._client.get_plans() matched: Optional[PlanSummary] = find_matching_plan( plans_response.plans, address ) if matched is None: raise ValueError(f"No MagicPlan found for address: {address!r}") raw_bytes: bytes = self._client.get_plan_raw(matched.id) magic_plan: MagicPlanPlan = MagicPlanPlan.model_validate( json.loads(raw_bytes)["data"] ) plan: Plan = map_plan(magic_plan) uploaded_file: UploadedFile = self._upload_raw_plan_json( plan_id=matched.id, raw_bytes=raw_bytes, uprn=uprn, hubspot_deal_id=request.hubspot_deal_id, ) with db_session() as session: save_plan(session, plan) session.add(uploaded_file) return plan def _upload_raw_plan_json( self, plan_id: str, raw_bytes: bytes, uprn: Optional[str], hubspot_deal_id: str, ) -> UploadedFile: compressed = gzip.compress(raw_bytes) if uprn is not None: s3_key = f"documents/uprn/{uprn}/magic_plan_{plan_id}.json.gz" else: s3_key = f"documents/hubspot_deal_id/{hubspot_deal_id}/magic_plan_{plan_id}.json.gz" save_data_to_s3(compressed, self._s3_bucket, s3_key) return UploadedFile( s3_file_bucket=self._s3_bucket, s3_file_key=s3_key, s3_upload_timestamp=datetime.now(timezone.utc), uprn=int(uprn) if uprn is not None else None, hubspot_deal_id=hubspot_deal_id, file_source=FileSourceEnum.MAGIC_PLAN.value, file_type=FileTypeEnum.MAGIC_PLAN_JSON.value, )