from __future__ import annotations from typing import Any, Optional, cast from sqlalchemy import delete, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlmodel import Session, col from domain.magicplan.models import ( Door, DoorVentilation, Floor, Plan, Room, Window, WindowVentilation, ) from infrastructure.postgres.magic_plan_tables import ( MagicPlanDoorModel, MagicPlanDoorVentilationModel, MagicPlanFloorModel, MagicPlanPlanModel, MagicPlanRoomModel, MagicPlanWindowModel, MagicPlanWindowVentilationModel, ) from repositories.magic_plan.magic_plan_repository import MagicPlanRepository class MagicPlanPostgresRepository(MagicPlanRepository): def __init__(self, session: Session) -> None: self._session = session def get_plan_by_uploaded_file_id(self, uploaded_file_id: int) -> Optional[Plan]: plan_row = self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanPlanModel).where( col(MagicPlanPlanModel.uploaded_file_id) == uploaded_file_id ) ).scalars().one_or_none() if plan_row is None: return None plan_id = cast(int, plan_row.id) floor_rows = list( self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanFloorModel).where( col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id ) ).scalars().all() ) floor_ids = [cast(int, f.id) for f in floor_rows] room_rows = list( self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanRoomModel).where( col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_ids) ) ).scalars().all() ) room_ids = [cast(int, r.id) for r in room_rows] window_rows = list( self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanWindowModel).where( col(MagicPlanWindowModel.magic_plan_room_id).in_(room_ids) ) ).scalars().all() ) window_ids = [cast(int, w.id) for w in window_rows] door_rows = list( self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanDoorModel).where( col(MagicPlanDoorModel.magic_plan_room_id).in_(room_ids) ) ).scalars().all() ) door_ids = [cast(int, d.id) for d in door_rows] window_vent_by_window_id = { cast(int, wv.magic_plan_window_id): wv for wv in self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanWindowVentilationModel).where( col(MagicPlanWindowVentilationModel.magic_plan_window_id).in_( window_ids ) ) ).scalars().all() } door_vent_by_door_id = { cast(int, dv.magic_plan_door_id): dv for dv in self._session.execute( # pyright: ignore[reportDeprecated] select(MagicPlanDoorVentilationModel).where( col(MagicPlanDoorVentilationModel.magic_plan_door_id).in_(door_ids) ) ).scalars().all() } windows_by_room_id: dict[int, list[Window]] = {} for w_row, w_id in zip(window_rows, window_ids): room_id = cast(int, w_row.magic_plan_room_id) wv_row = window_vent_by_window_id.get(w_id) ventilation = ( WindowVentilation( opening_type=wv_row.opening_type, num_openings=wv_row.num_openings, pct_openable=wv_row.pct_openable, trickle_vent_area_mm2=wv_row.trickle_vent_area_mm2, num_trickle_vents=wv_row.num_trickle_vents, ) if wv_row is not None else None ) windows_by_room_id.setdefault(room_id, []).append( Window( width_m=cast(float, w_row.width_m), height_m=cast(float, w_row.height_m), area_m2=cast(float, w_row.area_m2), ventilation=ventilation, ) ) doors_by_room_id: dict[int, list[Door]] = {} for d_row, d_id in zip(door_rows, door_ids): room_id = cast(int, d_row.magic_plan_room_id) dv_row = door_vent_by_door_id.get(d_id) ventilation = ( DoorVentilation(undercut_mm=dv_row.undercut_mm) if dv_row is not None else None ) doors_by_room_id.setdefault(room_id, []).append( Door( width_mm=cast(float, d_row.width_mm), height_mm=cast(float, d_row.height_mm), ventilation=ventilation, ) ) rooms_by_floor_id: dict[int, list[Room]] = {} for r_row, r_id in zip(room_rows, room_ids): floor_id = cast(int, r_row.magic_plan_floor_id) rooms_by_floor_id.setdefault(floor_id, []).append( Room( name=cast(str, r_row.name), width_m=cast(float, r_row.width_m), length_m=cast(float, r_row.length_m), area_m2=cast(float, r_row.area_m2), windows=windows_by_room_id.get(r_id, []), doors=doors_by_room_id.get(r_id, []), ) ) floors = [ Floor( level=f_row.level, name=None, rooms=rooms_by_floor_id.get(cast(int, f_row.id), []), ) for f_row in floor_rows ] return Plan( uid=cast(str, plan_row.magic_plan_uid), name=plan_row.name, address=plan_row.address, postcode=plan_row.postcode, floors=floors, ) def save(self, plan: Plan, uploaded_file_id: int) -> None: plan_id = self._upsert_plan(plan, uploaded_file_id) self._delete_children(plan_id) floor_ids = self._insert_floors(plan.floors, plan_id) room_ids = self._insert_rooms(plan.floors, floor_ids) window_ids, door_ids = self._insert_windows_and_doors(plan.floors, room_ids) self._insert_ventilation(plan.floors, window_ids, door_ids) def _upsert_plan(self, plan: Plan, uploaded_file_id: int) -> int: row_data: dict[str, Any] = MagicPlanPlanModel.from_domain( plan, uploaded_file_id ).model_dump(exclude={"id"}) stmt = ( pg_insert(MagicPlanPlanModel) .values(**row_data) .on_conflict_do_update( index_elements=["magic_plan_uid"], set_={k: v for k, v in row_data.items() if k != "magic_plan_uid"}, ) .returning(col(MagicPlanPlanModel.id)) ) return cast( int, self._session.execute(stmt).scalar_one() ) # pyright: ignore[reportDeprecated] def _delete_children(self, plan_id: int) -> None: floor_subq = ( select(col(MagicPlanFloorModel.id)) .where(col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id) .scalar_subquery() ) room_subq = ( select(col(MagicPlanRoomModel.id)) .where(col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq)) .scalar_subquery() ) window_subq = ( select(col(MagicPlanWindowModel.id)) .where(col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq)) .scalar_subquery() ) door_subq = ( select(col(MagicPlanDoorModel.id)) .where(col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq)) .scalar_subquery() ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanWindowVentilationModel).where( col(MagicPlanWindowVentilationModel.magic_plan_window_id).in_( window_subq ) ) ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanDoorVentilationModel).where( col(MagicPlanDoorVentilationModel.magic_plan_door_id).in_(door_subq) ) ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanWindowModel).where( col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq) ) ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanDoorModel).where( col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq) ) ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanRoomModel).where( col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq) ) ) self._session.execute( # pyright: ignore[reportDeprecated] delete(MagicPlanFloorModel).where( col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id ) ) def _insert_floors(self, floors: list[Floor], plan_id: int) -> list[int]: rows: list[dict[str, Any]] = [ MagicPlanFloorModel.from_domain(floor, plan_id).model_dump(exclude={"id"}) for floor in floors ] result = self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanFloorModel) .values(rows) .returning(col(MagicPlanFloorModel.id)) ) return cast(list[int], list(result.scalars().all())) def _insert_rooms(self, floors: list[Floor], floor_ids: list[int]) -> list[int]: rows: list[dict[str, Any]] = [ MagicPlanRoomModel.from_domain(room, floor_id).model_dump(exclude={"id"}) for floor, floor_id in zip(floors, floor_ids) for room in floor.rooms ] result = self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanRoomModel) .values(rows) .returning(col(MagicPlanRoomModel.id)) ) return cast(list[int], list(result.scalars().all())) def _insert_windows_and_doors( self, floors: list[Floor], room_ids: list[int] ) -> tuple[list[int], list[int]]: all_rooms = [room for floor in floors for room in floor.rooms] window_rows: list[dict[str, Any]] = [ MagicPlanWindowModel.from_domain(window, room_id).model_dump(exclude={"id"}) for room, room_id in zip(all_rooms, room_ids) for window in room.windows ] door_rows: list[dict[str, Any]] = [ MagicPlanDoorModel.from_domain(door, room_id).model_dump(exclude={"id"}) for room, room_id in zip(all_rooms, room_ids) for door in room.doors ] window_ids: list[int] = [] door_ids: list[int] = [] if window_rows: result = self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanWindowModel) .values(window_rows) .returning(col(MagicPlanWindowModel.id)) ) window_ids = cast(list[int], list(result.scalars().all())) if door_rows: result = self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanDoorModel) .values(door_rows) .returning(col(MagicPlanDoorModel.id)) ) door_ids = cast(list[int], list(result.scalars().all())) return window_ids, door_ids def _insert_ventilation( self, floors: list[Floor], window_ids: list[int], door_ids: list[int], ) -> None: all_rooms = [room for floor in floors for room in floor.rooms] all_windows = [w for room in all_rooms for w in room.windows] all_doors = [d for room in all_rooms for d in room.doors] window_vent_rows: list[dict[str, Any]] = [ MagicPlanWindowVentilationModel.from_domain(w.ventilation, wid).model_dump( exclude={"id"} ) for w, wid in zip(all_windows, window_ids) if w.ventilation is not None ] door_vent_rows: list[dict[str, Any]] = [ MagicPlanDoorVentilationModel.from_domain(d.ventilation, did).model_dump( exclude={"id"} ) for d, did in zip(all_doors, door_ids) if d.ventilation is not None ] if window_vent_rows: self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanWindowVentilationModel).values(window_vent_rows) ) if door_vent_rows: self._session.execute( # pyright: ignore[reportDeprecated] pg_insert(MagicPlanDoorVentilationModel).values(door_vent_rows) )