Model/orchestration/magic_plan_orchestrator.py

98 lines
3.3 KiB
Python

import gzip
import json
from datetime import datetime, timezone
import os
from typing import Optional, cast
from domain.magicplan.api.response import MagicPlanPlan, PlanSummary
from domain.magicplan.mapper import map_plan
from domain.magicplan.models import Plan
from backend.app.db.models.uploaded_file import (
FileSourceEnum,
FileTypeEnum,
UploadedFile,
)
from applications.magic_plan.address_matcher import find_matching_plan
from infrastructure.postgres.config import PostgresConfig
from infrastructure.postgres.engine import make_engine, make_session
from infrastructure.s3.s3_client import S3Client
from repositories.magic_plan.magic_plan_postgres_repository import (
MagicPlanPostgresRepository,
)
from infrastructure.magic_plan.magic_plan_client import MagicPlanClient
from applications.magic_plan.magic_plan_trigger_request import MagicPlanTriggerRequest
from utilities.logger import setup_logger
logger = setup_logger()
class MagicPlanOrchestrator:
def __init__(
self, magic_plan_api_client: MagicPlanClient, s3_client: S3Client
) -> None:
self._api_client = magic_plan_api_client
# self._s3_bucket = s3_bucket
self._s3_client = s3_client
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: list[PlanSummary] = self._api_client.get_plans()
matched: Optional[PlanSummary] = find_matching_plan(plans, address)
if matched is None:
raise ValueError(f"No MagicPlan found for address: {address!r}")
raw_bytes: bytes = self._api_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,
)
engine = make_engine(PostgresConfig.from_env(os.environ))
session = make_session(engine)
session.add(uploaded_file)
session.flush()
MagicPlanPostgresRepository(session).save(
plan, cast(int, uploaded_file.id)
) # TODO: refactor to use postgres Unit of Work
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"
self._s3_client.put_object(s3_key, compressed)
return UploadedFile(
s3_file_bucket=self._s3_client.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,
)