Add retry logic with tenacity

genotype_api/database/base_handler.py

@@ -1,10 +1,13 @@
 from dataclasses import dataclass
 from typing import Type
 
+from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
 from sqlalchemy.orm import DeclarativeBase, Query
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
+from genotype_api.config import settings
 from genotype_api.database.models import Analysis, Sample
 
 
@@ -15,6 +18,12 @@ class BaseHandler:
     def __init__(self, session: AsyncSession):
         self.session = session
 
+    @retry(
+        stop=stop_after_attempt(settings.max_retries),
+        wait=wait_fixed(settings.retry_delay),
+        retry=retry_if_exception_type(OperationalError),
+        reraise=True,
+    )
     def _get_query(self, table: Type[DeclarativeBase]) -> Query:
         """Return a query for the given table."""
         return select(table)

genotype_api/database/store.py

@@ -1,7 +1,10 @@
 """Module for the store handler."""
 
+from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
+from genotype_api.config import settings
 from genotype_api.database.crud.create import CreateHandler
 from genotype_api.database.crud.delete import DeleteHandler
 from genotype_api.database.crud.read import ReadHandler
@@ -24,6 +27,12 @@ def __init__(self, session: AsyncSession):
         UpdateHandler.__init__(self, session)
 
     @classmethod
+    @retry(
+        stop=stop_after_attempt(settings.max_retries),
+        wait=wait_fixed(settings.retry_delay),
+        retry=retry_if_exception_type(OperationalError),
+        reraise=True,
+    )
     async def create(cls) -> "Store":
         """Asynchronously create and return a Store instance with a session."""
         async with get_session() as session:  # Correctly use async context manager

pyproject.toml

@@ -35,6 +35,7 @@ uvicorn = "^0.29.0"
 uvloop = "^0.19.0"
 aiomysql = "^0.2.0"
 pytest-asyncio = "^0.24.0"
+tenacity = "^9.0.0"
 
 [tool.poetry.group.dev.dependencies]
 bump2version = "^1.0.1"