feat: implement uow

Author: fiqrikm18

src/modules/todo/application/create_todo/handler.py

@@ -15,10 +15,7 @@ async def execute(self, command: CreateTodoCommand, user_id: UUID) -> Todo:
         todo = Todo.create(
             title=command.title, user_id=user_id, description=command.description
         )
-        try:
+        async with self._unit_of_work:
             saved_todo = await self.todo_repo.save(todo)
             await self._unit_of_work.commit()
             return saved_todo
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise

src/modules/todo/application/delete_todo/handler.py

@@ -22,9 +22,6 @@ async def execute(self, todo_id: UUID, user_id: UUID) -> None:
                 "You do not have permission to delete this todo"
             )
 
-        try:
+        async with self._unit_of_work:
             await self.todo_repo.delete(todo_id)
             await self._unit_of_work.commit()
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise

src/modules/todo/application/update_todo/handler.py

@@ -40,10 +40,7 @@ async def execute(
             else:
                 todo.is_completed = False
 
-        try:
+        async with self._unit_of_work:
             saved_todo = await self.todo_repo.save(todo)
             await self._unit_of_work.commit()
             return saved_todo
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise

src/modules/user/application/login_user/handler.py

@@ -44,15 +44,12 @@ async def execute(self, command: LoginUserCommand) -> dict[str, str]:
             minutes=settings.REFRESH_TOKEN_EXPIRE_MINUTES
         )
 
-        try:
+        async with self._unit_of_work:
             new_rt = RefreshToken.create(
                 user_id=user.id, token_hash=token_hash, expires_at=expires_at
             )
             await self._refresh_token_repository.save(new_rt)
             await self._unit_of_work.commit()
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise
 
         return {
             "access_token": access_token,

src/modules/user/application/refresh_token/handler.py

@@ -35,7 +35,7 @@ async def execute(self, command: RefreshTokenCommand) -> dict:
         if stored_token.expires_at < datetime.now(timezone.utc):
             raise InvalidRefreshTokenError("Refresh token has expired")
 
-        try:
+        async with self._unit_of_work:
             stored_token.revoke()
             await self._refresh_token_repo.save(stored_token)
 
@@ -61,9 +61,6 @@ async def execute(self, command: RefreshTokenCommand) -> dict:
 
             await self._refresh_token_repo.save(new_refresh_token_entity)
             await self._unit_of_work.commit()
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise
 
         return {
             "access_token": new_access_token,

src/modules/user/application/register_user/handler.py

@@ -21,10 +21,7 @@ async def execute(self, command: RegisterUserCommand) -> User:
             command.email,
             password=hashed_password,
         )
-        try:
+        async with self._unit_of_work:
             saved_user = await self._user_repository.save(user=user)
             await self._unit_of_work.commit()
             return saved_user
-        except Exception:
-            await self._unit_of_work.rollback()
-            raise

src/shared/database/unit_of_work.py

@@ -1,3 +1,6 @@
+from types import TracebackType
+from typing import Self
+
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.shared.unit_of_work import UnitOfWork
@@ -6,9 +9,25 @@
 class SQLAlchemyUnitOfWork(UnitOfWork):
     def __init__(self, session: AsyncSession):
         self._session = session
+        self._committed = False
+
+    async def __aenter__(self) -> Self:
+        self._committed = False
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        if exc_type is not None or not self._committed:
+            await self.rollback()
+        return False
 
     async def commit(self) -> None:
         await self._session.commit()
+        self._committed = True
 
     async def rollback(self) -> None:
         await self._session.rollback()

src/shared/unit_of_work.py

@@ -1,7 +1,22 @@
 from abc import ABC, abstractmethod
+from types import TracebackType
+from typing import Self
 
 
 class UnitOfWork(ABC):
+    @abstractmethod
+    async def __aenter__(self) -> Self:
+        pass
+
+    @abstractmethod
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        pass
+
     @abstractmethod
     async def commit(self) -> None:
         pass

tests/shared/test_unit_of_work.py

@@ -37,3 +37,48 @@ async def run():
         assert session.rolled_back is True
 
     asyncio.run(run())
+
+
+def test_sqlalchemy_unit_of_work_rolls_back_when_scope_exits_without_commit():
+    async def run():
+        session = FakeSession()
+        unit_of_work = SQLAlchemyUnitOfWork(session)
+
+        async with unit_of_work:
+            pass
+
+        assert session.rolled_back is True
+        assert session.committed is False
+
+    asyncio.run(run())
+
+
+def test_sqlalchemy_unit_of_work_does_not_roll_back_after_commit():
+    async def run():
+        session = FakeSession()
+        unit_of_work = SQLAlchemyUnitOfWork(session)
+
+        async with unit_of_work:
+            await unit_of_work.commit()
+
+        assert session.committed is True
+        assert session.rolled_back is False
+
+    asyncio.run(run())
+
+
+def test_sqlalchemy_unit_of_work_rolls_back_and_propagates_scope_exception():
+    async def run():
+        session = FakeSession()
+        unit_of_work = SQLAlchemyUnitOfWork(session)
+
+        try:
+            async with unit_of_work:
+                raise RuntimeError("write failed")
+        except RuntimeError as exc:
+            assert str(exc) == "write failed"
+
+        assert session.rolled_back is True
+        assert session.committed is False
+
+    asyncio.run(run())

tests/todo/test_todo_flow.py

@@ -32,6 +32,14 @@ def __init__(self):
         self.committed = False
         self.rolled_back = False
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, traceback):
+        if exc_type is not None or not self.committed:
+            await self.rollback()
+        return False
+
     async def commit(self):
         self.committed = True