from __future__ import annotations import os from collections.abc import Callable from typing import Any, Optional, Protocol from sqlalchemy import Engine from sqlmodel import Session from applications.ara_first_run.ara_first_run_trigger_body import ( AraFirstRunTriggerBody, ) from domain.epc_prediction.epc_prediction import EpcPrediction from domain.property_baseline.calculator_rebaseliner import CalculatorRebaseliner from domain.sap10_calculator.calculator import Sap10Calculator from infrastructure.postgres.config import PostgresConfig from infrastructure.postgres.engine import make_engine from orchestration.property_baseline_orchestrator import PropertyBaselineOrchestrator from orchestration.ara_first_run_pipeline import AraFirstRunPipeline from orchestration.ingestion_orchestrator import ( ComparablesRepo, EpcFetcher, IngestionOrchestrator, PredictionAttributesReader, SolarFetcher, ) from orchestration.modelling_orchestrator import ModellingOrchestrator from orchestration.task_orchestrator import TaskOrchestrator from repositories.fuel_rates.fuel_rates_static_file_repository import ( FuelRatesStaticFileRepository, ) from repositories.geospatial.geospatial_repository import GeospatialRepository from repositories.postgres_unit_of_work import PostgresUnitOfWork from repositories.unit_of_work import UnitOfWork from utilities.aws_lambda.subtask_handler import subtask_handler # Module-scoped so the connection pool is reused across warm Lambda invocations # rather than rebuilt per invocation (ADR-0012). _engine: Optional[Engine] = None def _get_engine() -> Engine: global _engine if _engine is None: _engine = make_engine(PostgresConfig.from_env(dict(os.environ))) return _engine class _RunsFirstRun(Protocol): """The slice of AraFirstRunPipeline the handler delegates to.""" def run(self, command: AraFirstRunTriggerBody) -> None: ... def dispatch_first_run(body: dict[str, Any], *, pipeline: _RunsFirstRun) -> None: """Validate the raw event body and hand the command to the pipeline. The handler's entire decision logic — kept as a named seam so it is exercised without the Lambda runtime. No business logic: validate, delegate. """ trigger = AraFirstRunTriggerBody.model_validate(body) pipeline.run(trigger) def build_first_run_pipeline( *, unit_of_work: Callable[[], UnitOfWork], epc_fetcher: EpcFetcher, geospatial_repo: GeospatialRepository, solar_fetcher: SolarFetcher, comparables_repo: Optional[ComparablesRepo] = None, prediction_attributes_reader: Optional[PredictionAttributesReader] = None, ) -> AraFirstRunPipeline: """Compose the real three-stage pipeline on a Unit-of-Work factory. Each stage opens its own unit(s) and commits per batch (ADR-0012); the handler no longer holds a session. The source clients are passed in because their config is not settled — see ``_source_clients_from_env``. EPC Prediction gap-fill (ADR-0031) is the predictor itself (pure) plus two injected collaborators: the postcode-cohort source and the Landlord-Override attributes reader. Both default to None, so the feature is **off** until they are supplied — an EPC-less Property is then predicted into its predicted slot. The cohort repo is injected (not built here) because its EPC client is the same source client whose wiring is still pending; the attributes reader is the `property_overrides` read adapter built separately. Until both are passed, ingestion behaves exactly as before. """ return AraFirstRunPipeline( ingestion=IngestionOrchestrator( unit_of_work=unit_of_work, epc_fetcher=epc_fetcher, geospatial_repo=geospatial_repo, solar_fetcher=solar_fetcher, comparables_repo=comparables_repo, prediction_attributes_reader=prediction_attributes_reader, epc_prediction=EpcPrediction(), ), baseline=PropertyBaselineOrchestrator( unit_of_work=unit_of_work, # The calculator is load-bearing: effective=calculated for pre-10.2 # certs, lodged + divergence-logged at/above 10.2; a raise aborts the # batch (ADR-0013 amendment). rebaseliner=CalculatorRebaseliner(Sap10Calculator()), fuel_rates=FuelRatesStaticFileRepository(), ), modelling=ModellingOrchestrator( unit_of_work=unit_of_work, calculator=Sap10Calculator(), fuel_rates=FuelRatesStaticFileRepository(), ), ) def _source_clients_from_env() -> tuple[EpcFetcher, GeospatialRepository, SolarFetcher]: """The Ingestion source clients — EPC API, Google Solar, geospatial S3. TODO(deploy): their config (EPC auth token, Google Solar API key, geospatial S3 parquet reader), env-var names, and the pandas/s3fs runtime deps are not settled — that wiring is a separate Terraform piece, out of scope for #1136. Raises until then so the lambda fails loudly rather than half-running. """ raise NotImplementedError( "ara_first_run source-client wiring (EPC / Google Solar / geospatial) " "is pending the deploy/Terraform piece; see #1136." ) @subtask_handler() def handler( body: dict[str, Any], context: Any, task_orchestrator: TaskOrchestrator ) -> None: engine = _get_engine() unit_of_work: Callable[[], UnitOfWork] = lambda: PostgresUnitOfWork( lambda: Session(engine) ) epc_fetcher, geospatial_repo, solar_fetcher = _source_clients_from_env() pipeline = build_first_run_pipeline( unit_of_work=unit_of_work, epc_fetcher=epc_fetcher, geospatial_repo=geospatial_repo, solar_fetcher=solar_fetcher, ) dispatch_first_run(body, pipeline=pipeline)