"""SQS-triggered Lambda: fetch EPC (or predict) → run modelling → persist plan. One SQS message = one batch of properties sharing a portfolio, scenario, and (by caller convention) postcode. The handler reads ``property_ids``, ``portfolio_id``, ``scenario_id``, ``no_solar``, and ``dry_run`` from the message body, fetches or predicts each property's EPC, runs the full modelling pipeline (SAP10 → optimiser) via ``harness.console.run_modelling``, and persists the resulting Plan via ``PostgresUnitOfWork`` in one atomic transaction per property. When no lodged EPC is found, EPC Prediction (Path 3, ADR-0031) synthesises one from the postcode cohort. ``_cohort_cache`` is module-level so warm Lambda containers re-processing the same postcode avoid redundant fetches. All Measure Types are considered: pricing goes through ``catalogue_snapshot_with_off_catalogue_overrides`` so the measures the live ``material`` catalogue cannot supply (``secondary_heating_removal``, the glazing and heating gaps) are priced from the committed off-catalogue overlay instead of crashing. DB engine is module-scoped so the connection pool is reused across warm invocations (ADR-0012). The pool holds a single connection (``pool_size=1``): the handler reads everything up front — overrides, Scenario, a catalogue snapshot, and stored Solar — through one short-lived read Session, closes it, then writes each Property in a sequential Unit of Work, so the read and write Sessions never overlap. The orchestrator shares the same engine and releases its connection between bookkeeping commits, so one invocation uses one DB connection at a time. """ from __future__ import annotations import dataclasses import io import os from collections.abc import Callable, Generator from contextlib import contextmanager from typing import Any, Optional, cast import boto3 import pandas as pd # pyright: ignore[reportMissingTypeStubs] from sqlalchemy import Engine, text from sqlmodel import Session from datatypes.epc.domain.epc_property_data import ( BuildingPartIdentifier, EpcPropertyData, ) from domain.epc_prediction.comparable_properties import ( ComparableProperty, select_comparables, ) from domain.epc_prediction.epc_prediction import EpcPrediction from domain.epc_prediction.prediction_target import ( PredictionTarget, build_prediction_target, ) from domain.geospatial.coordinates import Coordinates from domain.geospatial.planning_restrictions import PlanningRestrictions from domain.geospatial.spatial_reference import SpatialReference from domain.property.property import Property, PropertyIdentity from domain.property_baseline.calculator_rebaseliner import CalculatorRebaseliner from domain.sap10_calculator.calculator import Sap10Calculator from domain.tasks.tasks import Source from harness.console import run_modelling from orchestration.task_orchestrator import TaskOrchestrator from orchestration.property_baseline_orchestrator import ( PropertyBaselineOrchestrator, ) from infrastructure.epc_client.epc_client_service import EpcClientService from infrastructure.postcodes_io.postcodes_io_client import PostcodesIoClient from infrastructure.postgres.config import PostgresConfig from infrastructure.postgres.engine import make_engine from infrastructure.solar.google_solar_api_client import ( BuildingInsightsNotFoundError, GoogleSolarApiClient, ) from applications.modelling_e2e.errors import ( DegeneratePredictionError, NoSameTypeComparablesError, UnresolvedPropertyTypeError, ) from applications.modelling_e2e.modelling_e2e_trigger_body import ( ModellingE2ETriggerBody, ) from repositories.comparable_properties.epc_comparable_properties_repository import ( EpcComparablePropertiesRepository, SkippedCohortCert, ) from repositories.geospatial.geospatial_s3_repository import ( GeospatialS3Repository, ParquetReader, ) from repositories.fuel_rates.fuel_rates_static_file_repository import ( FuelRatesStaticFileRepository, ) from repositories.postgres_unit_of_work import PostgresUnitOfWork from repositories.product.composite_product_repository import ( catalogue_snapshot_with_off_catalogue_overrides, ) from repositories.property.in_memory_property_overrides_reader import ( InMemoryPropertyOverridesReader, ) from repositories.property.landlord_override_overlays import overlays_from from repositories.property.override_backed_prediction_attributes_reader import ( OverrideBackedPredictionAttributesReader, ) from repositories.property.property_overrides_postgres_reader import ( PropertyOverridesPostgresReader, ) from repositories.property.property_overrides_reader import ( ResolvedPropertyOverrides, ) from repositories.scenario.scenario_postgres_repository import ( ScenarioPostgresRepository, ) from repositories.solar.solar_postgres_repository import SolarPostgresRepository from repositories.tasks.subtask_postgres_repository import ( SubTaskPostgresRepository, ) from repositories.tasks.task_postgres_repository import TaskPostgresRepository from utilities.aws_lambda.task_handler import task_handler from uuid import UUID from utilities.logger import setup_logger _engine: Optional[Engine] = None _cohort_cache: dict[str, list[ComparableProperty]] = {} # Broadened (nearby-postcode) cohorts, keyed by (seed postcode, target property # type): the early-stop walk depends on the type it is filling for, so two types # in the same postcode must not share a cached result. _nearby_cohort_cache: dict[tuple[str, str], list[ComparableProperty]] = {} logger = setup_logger() def _get_engine() -> Engine: global _engine if _engine is None: config = PostgresConfig.from_env(dict(os.environ)) # One connection per invocation: the handler reads everything up front # through one short-lived read Session, closes it, then writes each # Property in a sequential Unit of Work — so the read and write Sessions # never overlap and a single pooled connection suffices. The orchestrator # shares this engine (see ``_shared_engine_orchestrator``) and releases # its connection between bookkeeping commits, so it holds none during the # work. 32 concurrent containers × 1 connection = 32 against RDS. _engine = make_engine(dataclasses.replace(config, pool_size=1, max_overflow=0)) return _engine @contextmanager def _shared_engine_orchestrator() -> Generator[TaskOrchestrator, None, None]: """A ``TaskOrchestrator`` on the same module-scoped pooled engine as the modelling work — not a separate per-invocation NullPool engine. Its repositories commit on every ``save``/``create``, releasing the pooled connection between bookkeeping calls, so it holds none while the wrapped handler body runs. Combined with the read-then-write handler structure and ``pool_size=1``, the whole invocation uses one DB connection at a time.""" engine = _get_engine() with Session(engine) as session: yield TaskOrchestrator( task_repo=TaskPostgresRepository(session=session), subtask_repo=SubTaskPostgresRepository(session=session), ) def _s3_parquet_reader() -> ParquetReader: bucket = os.environ["DATA_BUCKET"] def read(key: str) -> pd.DataFrame: s3: Any = cast( Any, boto3.client("s3") ) # pyright: ignore[reportUnknownMemberType] raw = cast(bytes, s3.get_object(Bucket=bucket, Key=key)["Body"].read()) return pd.read_parquet(io.BytesIO(raw)) # type: ignore[return-value] return read def _spatial_for( geospatial: GeospatialS3Repository, uprn: int ) -> Optional[SpatialReference]: try: return geospatial.spatial_for(uprn) except Exception: # noqa: BLE001 return None def _solar_insights_for( solar_client: GoogleSolarApiClient, spatial: Optional[SpatialReference] ) -> Optional[dict[str, Any]]: if spatial is None or spatial.coordinates is None: return None try: return solar_client.get_building_insights( spatial.coordinates.longitude, spatial.coordinates.latitude ) except BuildingInsightsNotFoundError: return None def _dedupe_skipped( skipped: list[SkippedCohortCert], ) -> list[SkippedCohortCert]: """First occurrence of each skipped cert number (the same cert can appear in more than one postcode cohort across a batch).""" seen: set[str] = set() unique: list[SkippedCohortCert] = [] for cert in skipped: if cert.certificate_number not in seen: seen.add(cert.certificate_number) unique.append(cert) return unique def _predict_epc( *, property_id: int, uprn: int, postcode: str, portfolio_id: int, attributes_reader: OverrideBackedPredictionAttributesReader, coordinates: Optional[Coordinates], cohort_for: Callable[[str], list[ComparableProperty]], broaden: Callable[[PredictionTarget], list[ComparableProperty]], predictor: EpcPrediction, ) -> EpcPropertyData: """Synthesise an EpcPropertyData for an EPC-less property from its postcode cohort (EPC Prediction Path 3, ADR-0031). When the property's own postcode holds no same-type comparables (a sparse postcode — e.g. the only flat among houses), the cohort is broadened to the real unit postcodes physically nearest it (``broaden``) before giving up. Raises a specific ``PropertyNotModellableError`` subclass — naming the cause and carrying the property's identity — when it cannot predict: property_type unresolved, an empty same-type cohort, or a degenerate (no MAIN part) prediction. The per-property handler records ``str(exc)`` in the SubTask output, so the cause is debuggable from the output alone. """ attributes = attributes_reader.attributes_for(property_id) identity = PropertyIdentity( portfolio_id=portfolio_id, postcode=postcode, address="", uprn=uprn ) target = build_prediction_target(identity, coordinates, attributes) if target is None: raise UnresolvedPropertyTypeError( property_id=property_id, uprn=uprn, postcode=postcode, portfolio_id=portfolio_id, property_type=attributes.property_type, built_form=attributes.built_form, ) comparables = select_comparables(target, cohort_for(target.postcode)) broadened = False if not comparables.members: broadened = True comparables = select_comparables(target, broaden(target)) if not comparables.members: raise NoSameTypeComparablesError( property_id=property_id, uprn=uprn, postcode=postcode, portfolio_id=portfolio_id, property_type=target.property_type, broadened=broadened, ) predicted = predictor.predict(target, comparables) if not any( part.identifier is BuildingPartIdentifier.MAIN for part in predicted.sap_building_parts ): raise DegeneratePredictionError( property_id=property_id, uprn=uprn, postcode=postcode, portfolio_id=portfolio_id, property_type=target.property_type, cohort_size=len(comparables.members), ) return predicted @task_handler( task_source="modelling_e2e", source=Source.PROPERTY, orchestrator_cm=_shared_engine_orchestrator, pass_task_orchestrator=True, ) def handler(body: dict[str, Any], context: Any, orchestrator: TaskOrchestrator, task_id: UUID) -> None: trigger = ModellingE2ETriggerBody.model_validate(body) property_ids = trigger.property_ids portfolio_id = trigger.portfolio_id scenario_id = trigger.scenario_id no_solar = trigger.no_solar dry_run = trigger.dry_run logger.info( f"start property_ids={property_ids} portfolio={portfolio_id} " f"scenario={scenario_id} no_solar={no_solar} dry_run={dry_run}" ) engine = _get_engine() epc_client = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"]) geospatial = GeospatialS3Repository(_s3_parquet_reader()) solar_client = GoogleSolarApiClient(os.environ["GOOGLE_SOLAR_API_KEY"]) with engine.connect() as conn: uprn_rows = conn.execute( text("SELECT id, uprn FROM property WHERE id = ANY(:ids)"), {"ids": property_ids}, ).fetchall() postcode_rows = conn.execute( text("SELECT id, postcode FROM property WHERE id = ANY(:ids)"), {"ids": property_ids}, ).fetchall() uprns: dict[int, int] = {int(row[0]): int(row[1]) for row in uprn_rows} postcodes: dict[int, str] = {int(row[0]): (row[1] or "") for row in postcode_rows} # Pre-fetch every Property's overrides up front (each call opens and closes # its own short read Session) and serve them from memory through the loop, so # no override read Session is held open alongside a write Unit of Work. overrides_postgres_reader = PropertyOverridesPostgresReader(lambda: Session(engine)) overrides_by_pid: dict[int, ResolvedPropertyOverrides] = { pid: overrides_postgres_reader.overrides_for(pid) for pid in property_ids } overrides_reader = InMemoryPropertyOverridesReader(overrides_by_pid) prediction_attrs_reader = OverrideBackedPredictionAttributesReader(overrides_reader) comparables_repo = EpcComparablePropertiesRepository( epc_client, geospatial, nearby_postcodes=PostcodesIoClient() ) predictor = EpcPrediction() def _get_cohort(postcode: str) -> list[ComparableProperty]: if postcode not in _cohort_cache: _cohort_cache[postcode] = ( comparables_repo.candidates_for(postcode) if postcode else [] ) return _cohort_cache[postcode] def _broaden(target: PredictionTarget) -> list[ComparableProperty]: """The nearby-postcode cohort for a gated-out target — the real unit postcodes nearest it, walked until enough same-type comparables surface (ADR-0034). Memoised per (postcode, property_type) so co-located same-type misses share one walk.""" key = (target.postcode, target.property_type) if key not in _nearby_cohort_cache: _nearby_cohort_cache[key] = ( comparables_repo.candidates_near( target.postcode, target.coordinates, enough=lambda c: c.epc.property_type == target.property_type, ) if target.postcode else [] ) return _nearby_cohort_cache[key] # Re-establishes each lodged Property's Baseline Performance from the just- # persisted EPC (one UoW per property, committed after the Plan's). Predicted # Properties have no lodged figures, so they get no baseline (mirrors the e2e # runner and the ara_first_run Baseline stage). baseline_orchestrator = PropertyBaselineOrchestrator( unit_of_work=lambda: PostgresUnitOfWork(lambda: Session(engine)), rebaseliner=CalculatorRebaseliner(Sap10Calculator()), fuel_rates=FuelRatesStaticFileRepository(), ) read_session = Session(engine) try: # Read everything the write loop needs up front: the Scenario, an in-memory # snapshot of the catalogue (priced after the Session closes), and each # UPRN's stored Solar insights. Then close the read Session immediately so # its pooled connection is free before the loop — each Property's write # Unit of Work reuses that single connection rather than opening a second # alongside a held-open read Session. (The ``finally`` is the safety net.) scenario = ScenarioPostgresRepository(read_session).get_many([scenario_id])[0] products = catalogue_snapshot_with_off_catalogue_overrides(read_session) stored_solar: dict[int, Optional[dict[str, Any]]] = ( {} if no_solar else { uprn: SolarPostgresRepository(read_session).get(uprn) for uprn in set(uprns.values()) } ) read_session.close() for property_id in property_ids: child = orchestrator.create_child_subtask( task_id, inputs={"property_id": property_id} ) def _work(pid: int = property_id) -> None: uprn = uprns[pid] postcode = postcodes.get(pid, "") logger.info(f"property={pid} uprn={uprn} postcode={postcode!r}") spatial = _spatial_for(geospatial, uprn) restrictions = ( spatial.restrictions if spatial is not None else PlanningRestrictions() ) coordinates: Optional[Coordinates] = ( spatial.coordinates if spatial is not None else None ) epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn) overrides = overlays_from(overrides_reader.overrides_for(pid)) predicted_epc: Optional[EpcPropertyData] = None if epc is not None: logger.info(f"property={pid} lodged EPC found") effective_epc = Property( identity=PropertyIdentity( portfolio_id=portfolio_id, postcode=postcode, address="", uprn=uprn, ), epc=epc, landlord_overrides=overrides, ).effective_epc else: logger.info( f"property={pid} no lodged EPC — attempting prediction" ) predicted_epc = _predict_epc( property_id=pid, uprn=uprn, postcode=postcode, portfolio_id=portfolio_id, attributes_reader=prediction_attrs_reader, coordinates=coordinates, cohort_for=_get_cohort, broaden=_broaden, predictor=predictor, ) effective_epc = Property( identity=PropertyIdentity( portfolio_id=portfolio_id, postcode=postcode, address="", uprn=uprn, ), epc=None, predicted_epc=predicted_epc, landlord_overrides=overrides, ).effective_epc solar_insights: Optional[dict[str, Any]] solar_was_fetched = False if no_solar: solar_insights = None else: solar_insights = stored_solar.get(uprn) if solar_insights is None: solar_insights = _solar_insights_for(solar_client, spatial) solar_was_fetched = solar_insights is not None plan = run_modelling( effective_epc, planning_restrictions=restrictions, solar_insights=solar_insights, considered_measures=None, products=products, scenario=scenario, print_table=False, ) logger.info( f"property={pid} modelling complete " f"measures={len(plan.measures)}" ) if dry_run: measure_types = ( ", ".join(m.measure_type for m in plan.measures) or "none" ) logger.info( f"[dry_run] property={pid} " f"measures=[{measure_types}] — skipping DB write" ) return with PostgresUnitOfWork(lambda: Session(engine)) as uow: if epc is not None: uow.epc.save( epc, property_id=pid, portfolio_id=portfolio_id ) elif predicted_epc is not None: # Persist the synthesised EPC in the predicted slot (ADR-0031), # so the Baseline stage can re-hydrate it and downstream sees # the picture the Plan was modelled from. uow.epc.save( predicted_epc, property_id=pid, portfolio_id=portfolio_id, source="predicted", ) if spatial is not None: uow.spatial.save(uprn, spatial) if ( solar_was_fetched and solar_insights is not None and spatial is not None and spatial.coordinates is not None ): uow.solar.save( uprn, longitude=spatial.coordinates.longitude, latitude=spatial.coordinates.latitude, insights=solar_insights, ) uow.plan.save( plan, property_id=pid, scenario_id=scenario_id, portfolio_id=portfolio_id, is_default=scenario.is_default, ) uow.property.mark_modelled( pid, has_recommendations=bool(plan.measures) ) uow.commit() logger.info(f"property={pid} plan saved") baseline_orchestrator.run([pid]) logger.info(f"property={pid} baseline saved") try: orchestrator.run_subtask(child.id, work=_work) except Exception: # noqa: BLE001 pass # Cohort certs the mapper could not consume were skipped (not aborted on) # so prediction could proceed; surface them — with cert numbers — in the # subtask outputs so the mapper gaps can be closed later. skipped_certs: list[dict[str, str]] = [ {"certificate_number": s.certificate_number, "error": s.error} for s in _dedupe_skipped(comparables_repo.skipped) ] if skipped_certs: logger.info( f"skipped {len(skipped_certs)} unmappable cohort cert(s): " f"{[s['certificate_number'] for s in skipped_certs]}" ) finally: read_session.close()