from typing import Any, cast from sqlalchemy import delete, select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlmodel import Session, col from datatypes.magicplan.domain.models import Floor, Plan from backend.app.db.models.magic_plan import ( MagicPlanDoorModel, MagicPlanFloorModel, MagicPlanPlanModel, MagicPlanRoomModel, MagicPlanWindowModel, ) def save_plan(session: Session, plan: Plan, uploaded_file_id: int) -> None: plan_id: int = _upsert_plan(session, plan, uploaded_file_id) _delete_children(session, plan_id) floor_ids: list[int] = _insert_floors(session, plan.floors, plan_id) room_ids: list[int] = _insert_rooms(session, plan.floors, floor_ids) _insert_windows_and_doors(session, plan.floors, room_ids) def _upsert_plan(session: Session, plan: Plan, uploaded_file_id: int) -> int: stmt = ( pg_insert(MagicPlanPlanModel) .values( magic_plan_uid=plan.uid, name=plan.name, address=plan.address, postcode=plan.postcode, uploaded_file_id=uploaded_file_id, ) .on_conflict_do_update( index_elements=["magic_plan_uid"], set_={ "name": plan.name, "address": plan.address, "postcode": plan.postcode, "uploaded_file_id": uploaded_file_id, }, ) .returning(col(MagicPlanPlanModel.id)) ) row_id: int = session.execute(stmt).scalar_one() return row_id def _delete_children(session: Session, 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() ) session.execute( delete(MagicPlanWindowModel).where( col(MagicPlanWindowModel.magic_plan_room_id).in_(room_subq) ) ) session.execute( delete(MagicPlanDoorModel).where( col(MagicPlanDoorModel.magic_plan_room_id).in_(room_subq) ) ) session.execute( delete(MagicPlanRoomModel).where( col(MagicPlanRoomModel.magic_plan_floor_id).in_(floor_subq) ) ) session.execute( delete(MagicPlanFloorModel).where( col(MagicPlanFloorModel.magic_plan_plan_id) == plan_id ) ) def _insert_floors(session: Session, floors: list[Floor], plan_id: int) -> list[int]: rows: list[dict[str, Any]] = [ {"magic_plan_plan_id": plan_id, "level": floor.level} for floor in floors ] result = session.execute( pg_insert(MagicPlanFloorModel) .values(rows) .returning(col(MagicPlanFloorModel.id)) ) return cast(list[int], list(result.scalars().all())) def _insert_rooms( session: Session, floors: list[Floor], floor_ids: list[int] ) -> list[int]: rows: list[dict[str, Any]] = [ { "magic_plan_floor_id": floor_id, "name": room.name, "width_m": room.width_m, "length_m": room.length_m, "area_m2": room.area_m2, } for floor, floor_id in zip(floors, floor_ids) for room in floor.rooms ] result = session.execute( pg_insert(MagicPlanRoomModel).values(rows).returning(col(MagicPlanRoomModel.id)) ) return cast(list[int], list(result.scalars().all())) def _insert_windows_and_doors( session: Session, floors: list[Floor], room_ids: list[int] ) -> None: all_rooms = [room for floor in floors for room in floor.rooms] window_rows: list[dict[str, Any]] = [ { "magic_plan_room_id": room_id, "width_m": window.width_m, "height_m": window.height_m, "area_m2": window.area_m2, "opening_type": window.opening_type, } for room, room_id in zip(all_rooms, room_ids) for window in room.windows ] door_rows: list[dict[str, Any]] = [ { "magic_plan_room_id": room_id, "width_mm": door.width_mm, } for room, room_id in zip(all_rooms, room_ids) for door in room.doors ] if window_rows: session.execute(pg_insert(MagicPlanWindowModel).values(window_rows)) if door_rows: session.execute(pg_insert(MagicPlanDoorModel).values(door_rows))