import time from typing import List, Optional from sqlmodel import Session from utils.logger import setup_logger from backend.app.db.models.condition import ( AspectConditionModel, ElementModel, PropertyConditionSurveyModel, ) from backend.condition.domain.property_condition_survey import PropertyConditionSurvey from backend.app.db.connection import db_session logger = setup_logger() class ConditionPostgres: def bulk_insert_surveys( self, surveys: List[PropertyConditionSurvey], batch_size: Optional[int] = 100 ) -> None: logger.info( f"Preparing to load {len(surveys)} property surveys to Postgres. Mapping to SQLModel objects..." ) survey_models: List[PropertyConditionSurveyModel] = [ ConditionPostgres.map_survey_to_model(s) for s in surveys ] total: int = len(survey_models) logger.info( f"Finished mapping {total} surveys. Writing to database in batches of {batch_size}..." ) with db_session() as session: for start in range(0, total, batch_size): end = min(start + batch_size, total) batch = survey_models[start:end] t0: float = time.perf_counter() ConditionPostgres._insert_surveys_batch(batch, session) elapsed: float = time.perf_counter() - t0 logger.info( f"Inserted batch {start} - {end} ({len(batch)} surveys) in {elapsed} seconds", ) @staticmethod def map_survey_to_model( survey: PropertyConditionSurvey, ) -> PropertyConditionSurveyModel: survey_model = PropertyConditionSurveyModel( uprn=survey.uprn, date=survey.date, source=survey.source, elements=[], ) for element in survey.elements: element_model = ElementModel( element_type=element.element_type, element_instance=element.element_instance, aspect_conditions=[], ) for aspect in element.aspect_conditions: aspect_model = AspectConditionModel( aspect_type=aspect.aspect_type, aspect_instance=aspect.aspect_instance, value=aspect.value, quantity=aspect.quantity, install_date=aspect.install_date, renewal_year=aspect.renewal_year, comments=aspect.comments, ) element_model.aspect_conditions.append(aspect_model) survey_model.elements.append(element_model) return survey_model @staticmethod def _insert_surveys_batch( surveys: List[PropertyConditionSurveyModel], session: Session ) -> None: session.add_all(surveys) session.commit()