Model/backend/magic_plan/magic_plan_service.py

import gzip
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.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}")

        magic_plan: MagicPlanPlan = self._client.get_plan(matched.id)
        raw_bytes: bytes = self._client.get_plan_raw(matched.id)
        plan: Plan = map_plan(magic_plan)

        compressed = gzip.compress(raw_bytes)
        if uprn is not None:
            s3_key = f"documents/uprn/{uprn}/magic_plan_{matched.id}.json.gz"
        else:
            s3_key = f"documents/hubspot_deal_id/{request.hubspot_deal_id}/magic_plan_{matched.id}.json.gz"
        save_data_to_s3(compressed, self._s3_bucket, s3_key)

        with db_session() as session:
            save_plan(session, plan)

        return plan