feat: Implement `SELECT FOR UPDATE` and nested transactions to prevent race conditions in `OverhaulScope` and `MaintenanceType` operations.

main
Cizz22 4 months ago
parent f5c3cf774d
commit 6a73630d11

@ -56,7 +56,7 @@ async def update_scope(
scope_in: ScopeUpdate, scope_in: ScopeUpdate,
current_user: CurrentUser, 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: if not scope:
raise HTTPException( raise HTTPException(
@ -72,7 +72,7 @@ async def update_scope(
@router.post("/delete/{scope_id}", response_model=StandardResponse[ScopeRead]) @router.post("/delete/{scope_id}", response_model=StandardResponse[ScopeRead])
async def delete_scope(db_session: DbSession, scope_id: str): 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: if not scope:
raise HTTPException( raise HTTPException(

@ -20,10 +20,12 @@ from datetime import datetime
from uuid import UUID from uuid import UUID
async def get( async def get(
*, db_session: DbSession, overhaul_session_id: UUID *, db_session: DbSession, overhaul_session_id: UUID, for_update: bool = False
) -> Optional[OverhaulScope]: ) -> Optional[OverhaulScope]:
"""Returns a document based on the given document id.""" """Returns a document based on the given document id."""
query = Select(OverhaulScope).filter(OverhaulScope.id == overhaul_session_id).options(selectinload(OverhaulScope.maintenance_type)) 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) result = await db_session.execute(query)
return result.scalars().one_or_none() 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: if start_date and end_date:
duration_days = (end_date - start_date).days duration_days = (end_date - start_date).days
maintenance_type = Select(MaintenanceType).where(MaintenanceType.name == scope_in.type) # Prevent TOCTOU for maintenance type creation
maintenance_type = await db_session.execute(maintenance_type) maintenance_type = Select(MaintenanceType).where(MaintenanceType.name == scope_in.type).with_for_update()
maintenance_type = maintenance_type.scalars().one_or_none() maintenance_type_exec = await db_session.execute(maintenance_type)
maintenance_type = maintenance_type_exec.scalars().one_or_none()
if maintenance_type is None: if maintenance_type is None:
maintenance_type = MaintenanceType(name=scope_in.type) try:
db_session.add(maintenance_type) # Need to use nested transaction for safe concurrent creation
await db_session.flush() 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 # Create the OverhaulScope object with all hardcoded values
overhaul_session = OverhaulScope( overhaul_session = OverhaulScope(

Loading…
Cancel
Save