Store uploaded_file_id on magic_plan_plan row 🟥

backend/app/db/functions/magic_plan_functions.py

@@ -14,15 +14,15 @@ from backend.app.db.models.magic_plan import (
 )
 
 
-def save_plan(session: Session, plan: Plan) -> None:
-    plan_id: int = _upsert_plan(session, plan)
+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) -> int:
+def _upsert_plan(session: Session, plan: Plan, uploaded_file_id: int) -> int:
     stmt = (
         pg_insert(MagicPlanPlanModel)
         .values(
@@ -30,6 +30,7 @@ def _upsert_plan(session: Session, plan: Plan) -> int:
             name=plan.name,
             address=plan.address,
             postcode=plan.postcode,
+            uploaded_file_id=uploaded_file_id,
         )
         .on_conflict_do_update(
             index_elements=["magic_plan_uid"],
@@ -37,6 +38,7 @@ def _upsert_plan(session: Session, plan: Plan) -> int:
                 "name": plan.name,
                 "address": plan.address,
                 "postcode": plan.postcode,
+                "uploaded_file_id": uploaded_file_id,
             },
         )
         .returning(col(MagicPlanPlanModel.id))

backend/app/db/functions/tests/test_magic_plan_functions.py

@@ -36,7 +36,7 @@ def _count(session: Session, model: type[SQLModel]) -> int:
 
 def test_plan_row_present_after_save(db_session: Session, domain_plan: Plan) -> None:
     # Act
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
     # Assert
     assert _count(db_session, MagicPlanPlanModel) == 1
 
@@ -45,7 +45,7 @@ def test_floor_count_matches_domain(db_session: Session, domain_plan: Plan) -> N
     # Arrange
     expected = len(domain_plan.floors)
     # Act
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
     # Assert
     assert _count(db_session, MagicPlanFloorModel) == expected
 
@@ -54,7 +54,7 @@ def test_room_count_matches_domain(db_session: Session, domain_plan: Plan) -> No
     # Arrange
     expected = sum(len(f.rooms) for f in domain_plan.floors)
     # Act
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
     # Assert
     assert _count(db_session, MagicPlanRoomModel) == expected
 
@@ -63,7 +63,7 @@ def test_window_count_matches_domain(db_session: Session, domain_plan: Plan) ->
     # Arrange
     expected = sum(len(r.windows) for f in domain_plan.floors for r in f.rooms)
     # Act
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
     # Assert
     assert _count(db_session, MagicPlanWindowModel) == expected
 
@@ -72,15 +72,15 @@ def test_door_count_matches_domain(db_session: Session, domain_plan: Plan) -> No
     # Arrange
     expected = sum(len(r.doors) for f in domain_plan.floors for r in f.rooms)
     # Act
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
     # Assert
     assert _count(db_session, MagicPlanDoorModel) == expected
 
 
 def test_save_plan_idempotent(db_session: Session, domain_plan: Plan) -> None:
     # Act — call twice within the same session
-    save_plan(db_session, domain_plan)
-    save_plan(db_session, domain_plan)
+    save_plan(db_session, domain_plan, 1)
+    save_plan(db_session, domain_plan, 1)
     # Assert — same row counts as a single call
     assert _count(db_session, MagicPlanPlanModel) == 1
     assert _count(db_session, MagicPlanFloorModel) == len(domain_plan.floors)
@@ -93,3 +93,23 @@ def test_save_plan_idempotent(db_session: Session, domain_plan: Plan) -> None:
     assert _count(db_session, MagicPlanDoorModel) == sum(
         len(r.doors) for f in domain_plan.floors for r in f.rooms
     )
+
+
+def test_uploaded_file_id_stored_after_save(db_session: Session, domain_plan: Plan) -> None:
+    # Act
+    save_plan(db_session, domain_plan, 1)
+    # Assert
+    row = db_session.execute(select(MagicPlanPlanModel)).scalar_one()
+    assert row.uploaded_file_id == 1
+
+
+def test_save_plan_updates_uploaded_file_id_on_reingest(
+    db_session: Session, domain_plan: Plan
+) -> None:
+    # Arrange
+    save_plan(db_session, domain_plan, 1)
+    # Act
+    save_plan(db_session, domain_plan, 2)
+    # Assert
+    row = db_session.execute(select(MagicPlanPlanModel)).scalar_one()
+    assert row.uploaded_file_id == 2

backend/app/db/models/magic_plan.py

@@ -11,6 +11,7 @@ class MagicPlanPlanModel(SQLModel, table=True):
     name: Optional[str] = None
     address: Optional[str] = None
     postcode: Optional[str] = None
+    uploaded_file_id: Optional[int] = Field(default=None)
 
 
 class MagicPlanFloorModel(SQLModel, table=True):

backend/magic_plan/magic_plan_service.py

@@ -55,8 +55,9 @@ class MagicPlanService:
         )
 
         with db_session() as session:
-            save_plan(session, plan)
             session.add(uploaded_file)
+            session.flush()
+            save_plan(session, plan, uploaded_file.id)
 
         return plan