feat: implement uow

Author: fiqrikm18

src/core/database/session.py

@@ -1,13 +1,22 @@
+from collections.abc import AsyncGenerator
+
+from fastapi import Depends
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
 
 from src.core.config.setting import settings
+from src.shared.database.unit_of_work import SQLAlchemyUnitOfWork
+from src.shared.unit_of_work import UnitOfWork
 
 engine = create_async_engine(settings.DATABASE_URL, echo=False, future=True)
 AsyncSessionLocal = async_sessionmaker(
     engine, class_=AsyncSession, expire_on_commit=False, autoflush=False
 )
 
 
-async def get_db() -> AsyncSession:
+async def get_db() -> AsyncGenerator[AsyncSession]:
     async with AsyncSessionLocal() as session:
         yield session
+
+
+def get_unit_of_work(db: AsyncSession = Depends(get_db)) -> UnitOfWork:
+    return SQLAlchemyUnitOfWork(db)

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

@@ -3,14 +3,22 @@
 from src.modules.todo.application.create_todo.command import CreateTodoCommand
 from src.modules.todo.domain.entities.todo import Todo
 from src.modules.todo.domain.repositories.todo_repository import TodoRepository
+from src.shared.unit_of_work import UnitOfWork
 
 
 class CreateTodoHandler:
-    def __init__(self, todo_repo: TodoRepository):
+    def __init__(self, todo_repo: TodoRepository, unit_of_work: UnitOfWork):
         self.todo_repo = todo_repo
+        self._unit_of_work = unit_of_work
 
     async def execute(self, command: CreateTodoCommand, user_id: UUID) -> Todo:
         todo = Todo.create(
             title=command.title, user_id=user_id, description=command.description
         )
-        return await self.todo_repo.save(todo)
+        try:
+            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

@@ -5,11 +5,13 @@
     UnauthorizedTodoAccessError,
 )
 from src.modules.todo.domain.repositories.todo_repository import TodoRepository
+from src.shared.unit_of_work import UnitOfWork
 
 
 class DeleteTodoHandler:
-    def __init__(self, todo_repo: TodoRepository):
+    def __init__(self, todo_repo: TodoRepository, unit_of_work: UnitOfWork):
         self.todo_repo = todo_repo
+        self._unit_of_work = unit_of_work
 
     async def execute(self, todo_id: UUID, user_id: UUID) -> None:
         todo = await self.todo_repo.get_by_id(todo_id)
@@ -20,4 +22,9 @@ async def execute(self, todo_id: UUID, user_id: UUID) -> None:
                 "You do not have permission to delete this todo"
             )
 
-        await self.todo_repo.delete(todo_id)
+        try:
+            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

@@ -7,11 +7,13 @@
     UnauthorizedTodoAccessError,
 )
 from src.modules.todo.domain.repositories.todo_repository import TodoRepository
+from src.shared.unit_of_work import UnitOfWork
 
 
 class UpdateTodoHandler:
-    def __init__(self, todo_repo: TodoRepository):
+    def __init__(self, todo_repo: TodoRepository, unit_of_work: UnitOfWork):
         self.todo_repo = todo_repo
+        self._unit_of_work = unit_of_work
 
     async def execute(
         self, todo_id: UUID, command: UpdateTodoCommand, user_id: UUID
@@ -38,4 +40,10 @@ async def execute(
             else:
                 todo.is_completed = False
 
-        return await self.todo_repo.save(todo)
+        try:
+            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/infrastructure/repositories/todo_repository.py

@@ -50,7 +50,7 @@ async def save(self, todo: Todo) -> Todo:
             user_id=todo.user_id,
         )
         model = await self.db.merge(model)
-        await self.db.commit()
+        await self.db.flush()
         await self.db.refresh(model)
         return Todo(
             id=model.id,
@@ -62,4 +62,4 @@ async def save(self, todo: Todo) -> Todo:
 
     async def delete(self, todo_id: UUID) -> None:
         await self.db.execute(delete(TodoModel).where(TodoModel.id == todo_id))
-        await self.db.commit()
+        await self.db.flush()

src/modules/todo/presentation/dependency.py

@@ -1,6 +1,7 @@
 from fastapi import Depends
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from src.core.database.session import get_db, get_unit_of_work
 from src.modules.todo.application.delete_todo.handler import DeleteTodoHandler
 from src.modules.todo.application.create_todo.handler import CreateTodoHandler
 from src.modules.todo.application.list_todo.handler import GetTodosQueryHandler
@@ -9,7 +10,7 @@
 from src.modules.todo.infrastructure.repositories.todo_repository import (
     SQLAlchemyTodoRepository,
 )
-from src.core.database.session import get_db
+from src.shared.unit_of_work import UnitOfWork
 
 
 def get_todo_repository(db: AsyncSession = Depends(get_db)) -> TodoRepository:
@@ -18,20 +19,23 @@ def get_todo_repository(db: AsyncSession = Depends(get_db)) -> TodoRepository:
 
 def get_create_todo_handler(
     repo: TodoRepository = Depends(get_todo_repository),
+    unit_of_work: UnitOfWork = Depends(get_unit_of_work),
 ) -> CreateTodoHandler:
-    return CreateTodoHandler(repo)
+    return CreateTodoHandler(repo, unit_of_work)
 
 
 def get_update_todo_handler(
     repo: TodoRepository = Depends(get_todo_repository),
+    unit_of_work: UnitOfWork = Depends(get_unit_of_work),
 ) -> UpdateTodoHandler:
-    return UpdateTodoHandler(repo)
+    return UpdateTodoHandler(repo, unit_of_work)
 
 
 def get_delete_todo_handler(
     repo: TodoRepository = Depends(get_todo_repository),
+    unit_of_work: UnitOfWork = Depends(get_unit_of_work),
 ) -> DeleteTodoHandler:
-    return DeleteTodoHandler(repo)
+    return DeleteTodoHandler(repo, unit_of_work)
 
 
 def get_get_todos_query_handler(

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

@@ -12,16 +12,19 @@
 )
 from src.modules.user.domain.repositories.user_repository import UserRepository
 from src.shared.exceptions.credential_exception import InvalidCredentialsError
+from src.shared.unit_of_work import UnitOfWork
 
 
 class LoginUserCommandHandler:
     def __init__(
         self,
         user_repository: UserRepository,
         refresh_token_repository: RefreshTokenRepository,
+        unit_of_work: UnitOfWork,
     ):
         self._user_repository = user_repository
         self._refresh_token_repository = refresh_token_repository
+        self._unit_of_work = unit_of_work
 
     async def execute(self, command: LoginUserCommand) -> dict[str, str]:
         user = await self._user_repository.get_by_email(command.username)
@@ -41,10 +44,15 @@ async def execute(self, command: LoginUserCommand) -> dict[str, str]:
             minutes=settings.REFRESH_TOKEN_EXPIRE_MINUTES
         )
 
-        new_rt = RefreshToken.create(
-            user_id=user.id, token_hash=token_hash, expires_at=expires_at
-        )
-        await self._refresh_token_repository.save(new_rt)
+        try:
+            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

@@ -9,11 +9,17 @@
     RefreshTokenRepository,
 )
 from src.shared.exceptions.credential_exception import InvalidRefreshTokenError
+from src.shared.unit_of_work import UnitOfWork
 
 
 class RefreshTokenCommandHandler:
-    def __init__(self, refresh_token_repo: RefreshTokenRepository):
+    def __init__(
+        self,
+        refresh_token_repo: RefreshTokenRepository,
+        unit_of_work: UnitOfWork,
+    ):
         self._refresh_token_repo = refresh_token_repo
+        self._unit_of_work = unit_of_work
 
     def _hash_token(self, token: str) -> str:
         return hashlib.sha256(token.encode()).hexdigest()
@@ -29,30 +35,35 @@ async def execute(self, command: RefreshTokenCommand) -> dict:
         if stored_token.expires_at < datetime.now(timezone.utc):
             raise InvalidRefreshTokenError("Refresh token has expired")
 
-        stored_token.revoke()
-        await self._refresh_token_repo.save(stored_token)
+        try:
+            stored_token.revoke()
+            await self._refresh_token_repo.save(stored_token)
 
-        new_access_token = JWTService.create_access_token(
-            data={"sub": str(stored_token.user_id)}
-        )
+            new_access_token = JWTService.create_access_token(
+                data={"sub": str(stored_token.user_id)}
+            )
 
-        new_refresh_token_str = JWTService.create_refresh_token(
-            data={"sub": str(stored_token.user_id)}
-        )
+            new_refresh_token_str = JWTService.create_refresh_token(
+                data={"sub": str(stored_token.user_id)}
+            )
 
-        new_token_hash = self._hash_token(new_refresh_token_str)
+            new_token_hash = self._hash_token(new_refresh_token_str)
 
-        expires_at = datetime.now(timezone.utc) + timedelta(
-            minutes=settings.REFRESH_TOKEN_EXPIRE_MINUTES,
-        )
+            expires_at = datetime.now(timezone.utc) + timedelta(
+                minutes=settings.REFRESH_TOKEN_EXPIRE_MINUTES,
+            )
 
-        new_refresh_token_entity = RefreshToken.create(
-            user_id=stored_token.user_id,
-            token_hash=new_token_hash,
-            expires_at=expires_at,
-        )
+            new_refresh_token_entity = RefreshToken.create(
+                user_id=stored_token.user_id,
+                token_hash=new_token_hash,
+                expires_at=expires_at,
+            )
 
-        await self._refresh_token_repo.save(new_refresh_token_entity)
+            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

@@ -3,11 +3,13 @@
 from src.modules.user.domain.entities.user import User
 from src.modules.user.domain.exceptions.user_exception import UserAlreadyExistsError
 from src.modules.user.domain.repositories.user_repository import UserRepository
+from src.shared.unit_of_work import UnitOfWork
 
 
 class RegisterUserCommandHandler:
-    def __init__(self, user_repository: UserRepository):
+    def __init__(self, user_repository: UserRepository, unit_of_work: UnitOfWork):
         self._user_repository = user_repository
+        self._unit_of_work = unit_of_work
 
     async def execute(self, command: RegisterUserCommand) -> User:
         existing = await self._user_repository.get_by_email(command.email)
@@ -19,4 +21,10 @@ async def execute(self, command: RegisterUserCommand) -> User:
             command.email,
             password=hashed_password,
         )
-        return await self._user_repository.save(user=user)
+        try:
+            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/modules/user/infrastructure/repositories/refresh_token_repository.py

@@ -37,7 +37,7 @@ async def save(self, refresh_token: RefreshToken) -> RefreshToken:
             is_revoked=refresh_token.is_revoked,
         )
         model = await self.db.merge(model)
-        await self.db.commit()
+        await self.db.flush()
         await self.db.refresh(model)
         return RefreshToken(
             id=model.id,
@@ -56,4 +56,4 @@ async def revoke_by_user_id(self, user_id: UUID) -> None:
             )
             .values(is_revoked=True)
         )
-        await self.db.commit()
+        await self.db.flush()