[NOTICKET] Feat: retry_db decorator

externals/databases/pg_crud.py

@@ -1,7 +1,7 @@
 from typing import Optional, List
 from uuid import UUID
 
-from sqlalchemy import select, and_, update, delete, or_, func
+from sqlalchemy import select, and_, update, delete, func
 from sqlalchemy.ext.asyncio import AsyncSession
 from externals.databases.pg_models import (
     CVUser,
@@ -13,6 +13,10 @@ from externals.databases.pg_models import (
     CVMatching,
     CVScore,
 )
+from utils.decorator import retry_db
+
+
+
 
 # =========================
 # USER
@@ -20,6 +24,8 @@ from externals.databases.pg_models import (
 from externals.databases.schemas.user import UserCreate, UserResponse
 from utils.security import hash_password        
 
+
+@retry_db(retries=1, delay=2)
 async def get_user_by_username(
     db: AsyncSession,
     username: str,
@@ -30,6 +36,7 @@ async def get_user_by_username(
     return result.scalar_one_or_none()
 
 
+@retry_db(retries=1, delay=2)
 async def create_user(
     db: AsyncSession,
     user_in: UserCreate,
@@ -52,6 +59,8 @@ async def create_user(
 
     return user
 
+
+@retry_db(retries=1, delay=2)
 async def deactivate_user(
     db: AsyncSession,
     username: str,
@@ -71,6 +80,7 @@ async def deactivate_user(
     return user
 
 
+@retry_db(retries=1, delay=2)
 async def get_user_by_email(db, email: str):
     result = await db.execute(
         select(CVUser).where(CVUser.email == email)
@@ -78,6 +88,7 @@ async def get_user_by_email(db, email: str):
     return result.scalar_one_or_none()
 
 
+@retry_db(retries=1, delay=2)
 async def get_user_by_id(db, user_id: str):
     result = await db.execute(
         select(CVUser).where(CVUser.user_id == user_id)
@@ -90,6 +101,8 @@ async def get_user_by_id(db, user_id: str):
 # =========================
 from externals.databases.schemas.tenant import TenantCreate
 
+
+@retry_db(retries=1, delay=2)
 async def get_tenant_by_name(
     db: AsyncSession,
     tenant_name: str,
@@ -100,6 +113,7 @@ async def get_tenant_by_name(
     return result.scalar_one_or_none()
 
 
+@retry_db(retries=1, delay=2)
 async def create_tenant(
     db: AsyncSession,
     tenant_in: TenantCreate,
@@ -121,15 +135,7 @@ async def create_tenant(
 # FILE
 # =========================
 
-# async def get_file_by_filename(
-#     db: AsyncSession,
-#     filename: str,
-# ) -> Optional[CVFile]:
-#     stmt = select(CVFile).where(CVFile.filename == filename)
-#     result = await db.execute(stmt)
-#     return result.scalar_one_or_none()
-
-
+@retry_db(retries=1, delay=2)
 async def mark_file_extracted(
     db: AsyncSession,
     file_id: UUID,
@@ -166,6 +172,7 @@ async def create_cv_file(
     return cv_file
 
 
+@retry_db(retries=1, delay=2)
 async def delete_file_by_filename(
     db: AsyncSession,
     filename: str,

externals/databases/pg_models.py

@@ -3,6 +3,19 @@ from config.get_config import master_config
 from sqlalchemy import Column, String, Integer, TIMESTAMP, func
 from sqlalchemy.dialects.postgresql import UUID, ARRAY, DOUBLE_PRECISION, BOOLEAN
 from sqlalchemy.orm import declarative_base
+# from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
+# from config.constant import EnvPostgresConstants
+
+# engine = create_async_engine(
+#     EnvPostgresConstants.CONSTRING.replace('psycopg2', 'asyncpg'),
+#     pool_pre_ping=True,
+#     pool_recycle=300,
+# )
+
+# AsyncSessionLocal = async_sessionmaker(
+#     engine,
+#     expire_on_commit=False,
+# )
 
 
 Base = declarative_base()

utils/decorator.py

@@ -1,10 +1,11 @@
 import asyncio
+import asyncpg
 import inspect
-import functools
 import time
-import warnings
+
 from functools import wraps
 from typing import Callable, Any
+from sqlalchemy.exc import OperationalError, InterfaceError
 
 def trace_runtime(func: Callable) -> Callable:
     @wraps(func)
@@ -37,18 +38,43 @@ def trace_runtime(func: Callable) -> Callable:
         return sync_wrapper
     
 
-def deprecated(reason):
-    """
-    This is a decorator which can be used to mark functions
-    and classes as deprecated.
-    """
-    def decorator(func_or_class):
-        @functools.wraps(func_or_class)
-        def new_func(*args, **kwargs):
-            warnings.warn(f"Call to deprecated {func_or_class.__name__}."
-                          f" Reason: {reason}",
-                          category=FutureWarning,
-                          stacklevel=2)
-            return func_or_class(*args, **kwargs)
-        return new_func
+
+def retry_db(
+    retries: int = 3,
+    delay: float = 2.0,
+    backoff: float = 2.0,
+) -> Callable:
+
+    def decorator(func: Callable) -> Callable:
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs) -> Any:
+            current_delay = delay
+
+            for attempt in range(1, retries + 1):
+                try:
+                    return await func(*args, **kwargs)
+
+                except (
+                    OperationalError,
+                    InterfaceError,
+                    asyncpg.exceptions.PostgresConnectionError,
+                    asyncpg.exceptions.CannotConnectNowError,
+                    ConnectionError,
+                    TimeoutError,
+                ) as e:
+
+                    if attempt == retries:
+                        raise
+
+                    print(
+                        f"🔁 Retry {attempt}/{retries} for '{func.__name__}' "
+                        f"after {current_delay:.2f}s due to: {type(e).__name__}"
+                    )
+
+                    await asyncio.sleep(current_delay)
+                    current_delay *= backoff
+
+        return async_wrapper
+
     return decorator