Feature/improve uow (#14)

* add tenacity and upgrade libs;

* add a retry mechanism for UoW to handle SerializationError;

* allow config isolation level for relational database;

Author: EdwardPham1615

config.py

@@ -69,6 +69,9 @@ class RelationalDBConfig(BaseModel):
     url: str
     enable_log: bool
     enable_auto_migrate: bool
+    isolation_level: Optional[str] = (
+        "SERIALIZABLE"  # "READ UNCOMMITTED”, “READ COMMITTED”, “REPEATABLE READ”, “SERIALIZABLE”
+    )
 
 
 class CfgManagerConfig(BaseModel):

config/.env.example

@@ -15,6 +15,7 @@ RELATIONAL_DB__VENDOR=postgres
 RELATIONAL_DB__URL=postgresql+asyncpg://cleanarc:cleanarc!123@localhost:5432/cleanarc
 RELATIONAL_DB__ENABLE_LOG=true
 RELATIONAL_DB__ENABLE_AUTO_MIGRATE=true
+RELATIONAL_DB__ISOLATION_LEVEL='READ COMMITTED' # "READ UNCOMMITTED”, “READ COMMITTED”, “REPEATABLE READ”, “SERIALIZABLE”
 
 # === Authentication Service ===
 AUTHENTICATION_SERVICE__VENDOR=keycloak

internal/infrastructures/relational_db/patterns/unit_of_work.py

@@ -1,19 +1,30 @@
 import abc
+import logging
 from abc import abstractmethod
-from typing import Any, Callable, Optional, Type
+from typing import Callable, Optional
 
+from asyncpg.exceptions import SerializationError
 from sqlalchemy.ext.asyncio import (
     AsyncSession,
     AsyncSessionTransaction,
     async_scoped_session,
 )
+from tenacity import (
+    RetryError,
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from internal.infrastructures.relational_db import CommentRepo, PostRepo, UserRepo
 from internal.infrastructures.relational_db.abstraction import (
     AbstractCommentRepo,
     AbstractPostRepo,
     AbstractUserRepo,
 )
+from utils.logger_utils import get_shared_logger
 
 
 class AbstractUnitOfWork(abc.ABC):
@@ -59,30 +70,68 @@ def __init__(
         self._comment_repo_factory = comment_repo_factory
         self._user_repo_factory = user_repo_factory
 
-    async def __aenter__(self):
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=0.5, min=0.5, max=5),
+        retry=retry_if_exception_type(SerializationError),
+        before_sleep=before_sleep_log(get_shared_logger(), logging.WARNING),
+    )
+    async def _init_session(self):
+        """Initialize session with retry logic for serialization errors."""
         self._session = self._scoped_session_factory()
         self._transaction = await self._session.begin()
 
-        # register repo
+        # register repos
         self.post_repo = self._post_repo_factory(self._session)
         self.comment_repo = self._comment_repo_factory(self._session)
         self.user_repo = self._user_repo_factory(self._session)
         return self
 
-    async def __aexit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc: Optional[BaseException],
-        tb: Any,
-    ):
+    async def __aenter__(self):
+        try:
+            return await self._init_session()
+        except SerializationError as e:
+            # Ensure resources are cleaned up if all retries fail
+            if self._session:
+                await self._session.close()
+                await self._scoped_session_factory.remove()
+            raise
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=0.5, min=0.5, max=5),
+        retry=retry_if_exception_type(SerializationError),
+        before_sleep=before_sleep_log(get_shared_logger(), logging.WARNING),
+    )
+    async def _commit_transaction(self):
+        """Commit transaction with retry logic for serialization errors."""
+        try:
+            await self._transaction.commit()
+        except SerializationError:
+            # Rollback the failed transaction and start a new one for retry
+            await self._transaction.rollback()
+            self._transaction = await self._session.begin()
+            # Re-raise to trigger retry
+            raise
+
+    async def __aexit__(self, exc_type, exc, tb):
+        logger = get_shared_logger()
+
         try:
             if exc_type is None:
-                await self._transaction.commit()
+                try:
+                    # Try to commit with retries
+                    await self._commit_transaction()
+                except RetryError as e:
+                    # All retries exhausted
+                    logger.error(f"Max retries exceeded for commit: {e}")
+                    await self._transaction.rollback()
+                    raise SerializationError(
+                        "Failed to commit transaction after multiple retries"
+                    ) from e
             else:
+                # If there was an exception in the context, just rollback
                 await self._transaction.rollback()
-        except Exception:
-            await self._transaction.rollback()
-            raise
         finally:
             await self._session.close()
             await self._scoped_session_factory.remove()

internal/infrastructures/relational_db/postgres/__init__.py

@@ -19,15 +19,19 @@
 
 class PostgresDatabase:
     def __init__(
-        self, db_url: str, enable_log: bool = True, enable_migrations: bool = True
+        self,
+        db_url: str,
+        enable_log: bool = True,
+        enable_migrations: bool = True,
+        isolation_level: str = "SERIALIZABLE",
     ):
         self._db_url = db_url
         self._enable_log = enable_log
         self._enable_migrations = enable_migrations
         self._engine = create_async_engine(
             url=self._db_url,
             echo=self._enable_log,
-            isolation_level="SERIALIZABLE",
+            isolation_level=isolation_level,
         )
 
         self._session_factory = async_sessionmaker(

internal/patterns/dependency_injection.py

@@ -44,6 +44,7 @@ class Container(containers.DeclarativeContainer):
         db_url=config.relational_db.url,
         enable_log=config.relational_db.enable_log,
         enable_migrations=config.relational_db.enable_auto_migrate,
+        isolation_level=config.relational_db.isolation_level,
     )
 
     relational_db_scoped_session = providers.Resource(

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "ecs_logging",
     "pydantic-settings",
     "py-consul",
+    "tenacity",
 ]
 
 [tool.uv]