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(