diff --git a/infrastructure/postgres/magic_plan_tables.py b/infrastructure/postgres/magic_plan_tables.py index cbaf2dcc..abd06e1c 100644 --- a/infrastructure/postgres/magic_plan_tables.py +++ b/infrastructure/postgres/magic_plan_tables.py @@ -108,12 +108,24 @@ class MagicPlanWindowVentilationModel(SQLModel, table=True): id: Optional[int] = Field(default=None, primary_key=True) magic_plan_window_id: int = Field(foreign_key="magic_plan_window.id") + opening_type: Optional[str] = None + num_openings: Optional[int] = None + pct_openable: Optional[int] = None + trickle_vent_area_mm2: Optional[int] = None + num_trickle_vents: Optional[int] = None @classmethod def from_domain( - cls, _ventilation: WindowVentilation, _window_id: int + cls, ventilation: WindowVentilation, window_id: int ) -> "MagicPlanWindowVentilationModel": - raise NotImplementedError + return cls( + magic_plan_window_id=window_id, + opening_type=ventilation.opening_type, + num_openings=ventilation.num_openings, + pct_openable=ventilation.pct_openable, + trickle_vent_area_mm2=ventilation.trickle_vent_area_mm2, + num_trickle_vents=ventilation.num_trickle_vents, + ) class MagicPlanDoorVentilationModel(SQLModel, table=True): @@ -121,9 +133,13 @@ class MagicPlanDoorVentilationModel(SQLModel, table=True): id: Optional[int] = Field(default=None, primary_key=True) magic_plan_door_id: int = Field(foreign_key="magic_plan_door.id") + undercut_mm: Optional[float] = None @classmethod def from_domain( - cls, _ventilation: DoorVentilation, _door_id: int + cls, ventilation: DoorVentilation, door_id: int ) -> "MagicPlanDoorVentilationModel": - raise NotImplementedError + return cls( + magic_plan_door_id=door_id, + undercut_mm=ventilation.undercut_mm, + ) diff --git a/repositories/magic_plan/magic_plan_postgres_repository.py b/repositories/magic_plan/magic_plan_postgres_repository.py index 3c464b24..57d9ff25 100644 --- a/repositories/magic_plan/magic_plan_postgres_repository.py +++ b/repositories/magic_plan/magic_plan_postgres_repository.py @@ -9,10 +9,12 @@ from sqlmodel import Session, col from datatypes.magicplan.domain.models import Floor, Plan from infrastructure.postgres.magic_plan_tables import ( MagicPlanDoorModel, + MagicPlanDoorVentilationModel, MagicPlanFloorModel, MagicPlanPlanModel, MagicPlanRoomModel, MagicPlanWindowModel, + MagicPlanWindowVentilationModel, ) from repositories.magic_plan.magic_plan_repository import MagicPlanRepository @@ -26,7 +28,8 @@ class MagicPlanPostgresRepository(MagicPlanRepository): self._delete_children(plan_id) floor_ids = self._insert_floors(plan.floors, plan_id) room_ids = self._insert_rooms(plan.floors, floor_ids) - self._insert_windows_and_doors(plan.floors, room_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( @@ -54,6 +57,26 @@ class MagicPlanPostgresRepository(MagicPlanRepository): .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) @@ -102,7 +125,7 @@ class MagicPlanPostgresRepository(MagicPlanRepository): def _insert_windows_and_doors( self, floors: list[Floor], room_ids: list[int] - ) -> None: + ) -> 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( @@ -116,7 +139,53 @@ class MagicPlanPostgresRepository(MagicPlanRepository): 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: - self._session.execute(pg_insert(MagicPlanWindowModel).values(window_rows)) # pyright: ignore[reportDeprecated] + 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: - self._session.execute(pg_insert(MagicPlanDoorModel).values(door_rows)) # pyright: ignore[reportDeprecated] + 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) + ) diff --git a/tests/repositories/magic_plan/test_magic_plan_postgres_repository.py b/tests/repositories/magic_plan/test_magic_plan_postgres_repository.py index a71595db..d19896f2 100644 --- a/tests/repositories/magic_plan/test_magic_plan_postgres_repository.py +++ b/tests/repositories/magic_plan/test_magic_plan_postgres_repository.py @@ -110,3 +110,5 @@ def test_save_is_idempotent(db_engine: Engine) -> None: assert len(session.exec(select(MagicPlanRoomModel)).all()) == 1 assert len(session.exec(select(MagicPlanWindowModel)).all()) == 1 assert len(session.exec(select(MagicPlanDoorModel)).all()) == 1 + assert len(session.exec(select(MagicPlanWindowVentilationModel)).all()) == 1 + assert len(session.exec(select(MagicPlanDoorVentilationModel)).all()) == 1