From 6a73630d11e68c50fc460078bea02be89a4cbf4b Mon Sep 17 00:00:00 2001
From: Cizz22 <cisatraa@gmail.com>
Date: Mon, 9 Mar 2026 12:28:28 +0700
Subject: [PATCH] feat: Implement `SELECT FOR UPDATE` and nested transactions
 to prevent race conditions in `OverhaulScope` and `MaintenanceType`
 operations.

---
 src/overhaul_scope/router.py  |  4 ++--
 src/overhaul_scope/service.py | 24 +++++++++++++++++-------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/overhaul_scope/router.py b/src/overhaul_scope/router.py
index 8981e64..9a069c6 100644
--- a/src/overhaul_scope/router.py
+++ b/src/overhaul_scope/router.py
@@ -56,7 +56,7 @@ async def update_scope(
     scope_in: ScopeUpdate,
     current_user: CurrentUser,
 ):
-    scope = await get(db_session=db_session, scope_id=scope_id)
+    scope = await get(db_session=db_session, overhaul_session_id=scope_id, for_update=True)
 
     if not scope:
         raise HTTPException(
@@ -72,7 +72,7 @@ async def update_scope(
 
 @router.post("/delete/{scope_id}", response_model=StandardResponse[ScopeRead])
 async def delete_scope(db_session: DbSession, scope_id: str):
-    scope = await get(db_session=db_session, scope_id=scope_id)
+    scope = await get(db_session=db_session, overhaul_session_id=scope_id, for_update=True)
 
     if not scope:
         raise HTTPException(
diff --git a/src/overhaul_scope/service.py b/src/overhaul_scope/service.py
index 47df892..f24f861 100644
--- a/src/overhaul_scope/service.py
+++ b/src/overhaul_scope/service.py
@@ -20,10 +20,12 @@ from datetime import datetime
 from uuid import UUID
 
 async def get(
-    *, db_session: DbSession, overhaul_session_id: UUID
+    *, db_session: DbSession, overhaul_session_id: UUID, for_update: bool = False
 ) -> Optional[OverhaulScope]:
     """Returns a document based on the given document id."""
     query = Select(OverhaulScope).filter(OverhaulScope.id == overhaul_session_id).options(selectinload(OverhaulScope.maintenance_type))
+    if for_update:
+        query = query.with_for_update()
     result = await db_session.execute(query)
     return result.scalars().one_or_none()
 
@@ -71,14 +73,22 @@ async def create(*, db_session: DbSession, scope_in: ScopeCreate):
     if start_date and end_date:
         duration_days = (end_date - start_date).days
 
-    maintenance_type = Select(MaintenanceType).where(MaintenanceType.name == scope_in.type)
-    maintenance_type = await db_session.execute(maintenance_type)
-    maintenance_type = maintenance_type.scalars().one_or_none()
+    # Prevent TOCTOU for maintenance type creation
+    maintenance_type = Select(MaintenanceType).where(MaintenanceType.name == scope_in.type).with_for_update()
+    maintenance_type_exec = await db_session.execute(maintenance_type)
+    maintenance_type = maintenance_type_exec.scalars().one_or_none()
 
     if maintenance_type is None:
-        maintenance_type = MaintenanceType(name=scope_in.type)
-        db_session.add(maintenance_type)
-        await db_session.flush()
+        try:
+            # Need to use nested transaction for safe concurrent creation
+            async with db_session.begin_nested():
+                maintenance_type = MaintenanceType(name=scope_in.type)
+                db_session.add(maintenance_type)
+                await db_session.flush()
+        except BaseException:
+            # Another concurrent request already created it
+            maintenance_type_exec = await db_session.execute(Select(MaintenanceType).where(MaintenanceType.name == scope_in.type).with_for_update())
+            maintenance_type = maintenance_type_exec.scalars().one_or_none()
 
     # Create the OverhaulScope object with all hardcoded values
     overhaul_session = OverhaulScope(